"""
4.	利用Tensorflow或Pytorch深度学习框架，建立由适当卷积层、池化层和全连接构成的CNN模型，进行“验证码”识别。验证码数据集参考下图
"""

import os
import cv2 as cv
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import losses, metrics, optimizers, layers, activations, callbacks
from sklearn.model_selection import train_test_split

np.random.seed(777)
tf.random.set_seed(777)

# 按下面要求完成CNN模型训练，并评估“验证码”识别的准确率。（15分）
# ①	定义读取数据函数read_data()，划分数据集为训练集（0.9）和测试集（0.1）


def read_data(dir):
    x = []
    y = []
    names = os.listdir(dir)
    for name in names:
        name_split = os.path.splitext(name)

        # only for jpg files
        ext = name_split[1].lower()
        if '.jpg' != ext:
            continue

        # label
        label_str = name_split[0]
        label = [int(ch) for ch in list(label_str)]
        y.append(label)

        path = os.path.join(dir, name)
        img = cv.imread(path, cv.IMREAD_COLOR)
        x.append(img)

    x = np.float32(x) / 255.
    y = np.int32(y)
    eye = np.eye(10)
    y = eye[y]
    y = y.reshape(-1, 40)
    return x, y


dir = 'data/vcode_data'
x, y = read_data(dir)
print('x', x.shape)
print('y', y.shape)
# split
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.9, random_state=777)

# ②	定义cnn模型，并用训练集数据训练模型


def ConvBnRelu(ch, ksize):
    return keras.Sequential([
        layers.Conv2D(ch, ksize, (1, 1), 'same'),
        layers.BatchNormalization(),
        layers.ReLU()
    ])


# 定义cnn模型
inputs = keras.Input(shape=(x_train.shape[1], x_train.shape[2], x_train.shape[3]))
x = ConvBnRelu(16, (3, 3))(inputs)
x = layers.AvgPool2D((3, 3), (2, 2), 'same')(x)
x = ConvBnRelu(16, (3, 3))(x)
x = layers.AvgPool2D((3, 3), (2, 2), 'same')(x)
x = ConvBnRelu(32, (3, 3))(x)
x = layers.AvgPool2D((3, 3), (2, 2), 'same')(x)
x = ConvBnRelu(32, (3, 3))(x)
x = layers.AvgPool2D((3, 3), (2, 2), 'same')(x)
x = layers.Flatten()(x)
x = layers.Dense(256, activation=activations.relu)(x)
x = layers.Dense(64, activation=activations.relu)(x)
x = layers.Dense(40, activation=activations.sigmoid)(x)
model = keras.Model(inputs, x)
model.summary()


def my_acc(y_true, y_pred):
    """符合题意的自定义准确率"""
    y_true = tf.reshape(y_true, (-1, 4, 10))
    y_true_ints = tf.argmax(y_true, axis=2)
    y_pred = tf.reshape(y_pred, (-1, 4, 10))
    y_pred_ints = tf.argmax(y_pred, axis=2)
    eq = y_true_ints == y_pred_ints
    eq = tf.reduce_all(eq, axis=1)
    eq = tf.cast(eq, dtype=tf.float32)
    eq_mean = tf.reduce_mean(eq)
    return eq_mean


model.compile(
    loss=losses.binary_crossentropy,
    optimizer=optimizers.Adam(learning_rate=0.0001),
    metrics=my_acc
)

# 用训练集数据训练模型
BATCH_SIZE = 64
N_EPOCHS = 5  # 为了快速演示，这里设为5，正式代码要适当加大
model.fit(x_train, y_train,
          batch_size=BATCH_SIZE, epochs=N_EPOCHS,
          validation_data=(x_test, y_test), validation_batch_size=BATCH_SIZE)

# ③	用测试集数据评估模型“验证码”识别的准确率
res = model.evaluate(x_test, y_test, batch_size=BATCH_SIZE)
print('注意：由于使用了符合题意的自定义准确率，且样本过少，测试准确率可能为0')
print(f'用测试集数据评估模型“验证码”识别的准确率: {res[1]}')